import numpy as np
import pandas as pd

from datetime import datetime

# Visualization libraries
import seaborn as sns
import matplotlib
import matplotlib.pyplot as plt
import plotly.express as px

# ML/DL libraries
import tensorflow as tf
from tensorflow.keras import layers
import tensorflow.keras.applications as models
import tensorflow_ranking as tfr
import tensorflow_recommenders as tfrs
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import train_test_split
import keras_tuner as kt

# Misc
import math
import random
from joblib import Parallel, delayed
import PIL
from PIL import Image
import os
import pathlib
import warnings
from pprint import pprint
from numba import cuda 
from tqdm import tqdm

# Local
from metrics import dcg
from metrics import ndcg
from metrics import permutation_accuracy
from metrics import accuracy_by_rank
from metrics import accuracy_at_k
from metrics import batch_kendalls_tau
from metrics import batch_spearmans_rho
from metrics import ordered_pair_accuracy

from constants import BATCH_SIZE
from constants import DATASET_DIRECTORY
from constants import IMAGE_TYPE
from constants import IMAGENET_INPUT_H
from constants import IMAGENET_INPUT_W
from constants import IMAGENET_INPUT_D
from constants import IMAGENET_MAX_PIXEL
from constants import IMAGENET_IMAGE_SIZE
from constants import IMAGENET_SHAPE
from constants import NUM_CLASSES


def image_generator(data_dir):
    data_dir = pathlib.Path(data_dir)
    subdirs = list(data_dir.iterdir())

    for subdir in subdirs:
        if subdir.is_dir():
            subdir_images = []
            subdir_labels = []

            for i in range(1, NUM_CLASSES + 1):
                filepath = subdir / f'{i}.png'
                img_raw = tf.io.read_file(str(filepath))
                img = tf.image.decode_png(img_raw, channels=IMAGENET_INPUT_D)
                img = tf.image.resize(img, [IMAGENET_INPUT_H, IMAGENET_INPUT_W])
                img /= IMAGENET_MAX_PIXEL
                label = NUM_CLASSES - i  
                
                subdir_images.append(img)
                subdir_labels.append(label)
            
            paired = list(zip(subdir_images, subdir_labels))
            np.random.shuffle(paired)
            images_, labels_ = zip(*paired)

            yield tf.stack(images_), tf.stack(labels_)


def create_dataset(data_dir, batch_size=BATCH_SIZE):
    dataset = tf.data.Dataset.from_generator(
        lambda: image_generator(data_dir),
        output_signature=(
            tf.TensorSpec(shape=(NUM_CLASSES, IMAGENET_INPUT_H, IMAGENET_INPUT_W, IMAGENET_INPUT_D), dtype=tf.float32),
            tf.TensorSpec(shape=(NUM_CLASSES,), dtype=tf.int32)
        )
    )
    
    dataset = dataset.batch(batch_size).prefetch(tf.data.experimental.AUTOTUNE)

    dataset_size = sum(1 for _ in dataset)
    train_size = int(dataset_size * 0.7)
    val_size = int(dataset_size * 0.15)
   
    train_dataset = dataset.take(train_size)
    test_dataset = dataset.skip(train_size)
    val_dataset = test_dataset.skip(val_size)
    test_dataset = test_dataset.take(val_size)
    
    return train_dataset, test_dataset, val_dataset


def save_model_with_timestamp_and_lr(model, lr, freezed="unfreezed", base_dir="tesla_saved_models"):
    class_name = model.__class__.__name__
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    model_name = f"{class_name}_{timestamp}_{freezed}_{str(lr)}"
    model_dir = os.path.join(base_dir, model_name)
    
    if not os.path.exists(base_dir):
        os.makedirs(base_dir)
    model.save(model_dir)

    print(f"Model saved in {model_dir} as {model_name}")


def predict_rank(model, images):
    scores = model(images, training=False)
    ranks = tf.argsort(tf.argsort(scores, axis=1), axis=1)
    return ranks


def compute_dataset_metrics(model, dataset, max_k=5):
    aggregated_metrics = {
        'NDCG': 0.0,
        'permutation_accuracy': 0.0,
        'Kendalls Tau': 0.0,
        'Spearmans Rho': 0.0,
        'OPA': 0.0,
        'accuracy_by_rank': np.zeros(max_k),
        'accuracy@k': np.zeros(max_k)
    }
    total_batches = 0
    
    for images, labels in tqdm(dataset, desc="Processing batches"):
        scores = model.predict(images)
        ranks = np.argsort(np.argsort(scores, axis=1), axis=1)
        
        aggregated_metrics['OPA'] += ordered_pair_accuracy(labels, ranks)
        aggregated_metrics['Kendalls Tau'] += batch_kendalls_tau(labels, ranks)
        aggregated_metrics['Spearmans Rho'] += batch_spearmans_rho(labels, ranks)
        aggregated_metrics['NDCG'] += ndcg(labels, ranks)
        aggregated_metrics['permutation_accuracy'] += permutation_accuracy(labels, ranks)
        aggregated_metrics['accuracy@k'] += np.array(accuracy_at_k(labels, ranks))

        accuracy_by_rank_ = []
        for k in range(0, max_k):
            accuracy_by_rank_.append(accuracy_by_rank(labels, ranks, k))

        aggregated_metrics['accuracy_by_rank'] += np.array(accuracy_by_rank_)

        total_batches += 1
    
    aggregated_metrics = {k: v / total_batches for k, v in aggregated_metrics.items()}

    return aggregated_metrics


def plot_metrics(metrics):    
    ndcg = metrics['NDCG']
    permutation_accuracy = metrics['permutation_accuracy']
    accuracy_by_rank = metrics['accuracy_by_rank']
    accuracy_at_k = metrics['accuracy@k']
    
    max_k = len(accuracy_by_rank)

    ranks = list(range(1, max_k + 1))[::-1]

    fig, axs = plt.subplots(1, 2, figsize=(14, 6))

    axs[0].bar(ranks, accuracy_by_rank, color='skyblue')
    axs[0].set_title('Accuracy by Image')
    axs[0].set_xlabel('Image')
    axs[0].set_ylabel('Accuracy')
    axs[0].set_xticks(ranks)
    axs[0].set_xticklabels(ranks)
    axs[0].set_ylim(0, 1.1)

    axs[1].bar(ranks, accuracy_at_k[::-1], color='lightgreen')
    axs[1].set_title('Accuracy@k')
    axs[1].set_xlabel('k')
    axs[1].set_ylabel('Accuracy')
    axs[1].set_xticks(ranks)
    axs[1].set_xticklabels(ranks)
    axs[1].set_ylim(0, 1.1)

    textstr = '\n'.join((
        f'NDCG: {ndcg:.2f}',
        f'Permutation Accuracy: {permutation_accuracy:.2f}'
    ))

    props = dict(boxstyle='round', facecolor='wheat', alpha=0.5)
    fig.text(0.5, 0.95, textstr, fontsize=12, verticalalignment='top', bbox=props, ha='center')

    plt.tight_layout(rect=[0, 0, 1, 0.93])
    plt.show()